project(mnist)

# Torch
find_package(Torch REQUIRED PATHS $ENV{Torch_DIR} HINTS
             $ENV{HOME}/libtorch/share/cmake/Torch)
# Download the MNIST dataset
find_package(
  Python3
  COMPONENTS Interpreter
  REQUIRED)
option(DOWNLOAD_MNIST "Download the MNIST dataset from the internet" ON)
if(DOWNLOAD_MNIST)
  message(STATUS "Downloading MNIST dataset")
  # cmake-format: off
               execute_process(
                 COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tools/download_mnist.py 
                         -d ${CMAKE_BINARY_DIR}/data 
                         ERROR_VARIABLE 
                         DOWNLOAD_ERROR)
               # cmake-format: on
  if(DOWNLOAD_ERROR)
    message(FATAL_ERROR "Error downloading MNIST dataset: ${DOWNLOAD_ERROR}")
  endif()
endif()

# MPI_BGV_64
# create_target(
#   test_mpi_bgv_mnist_64
#   "test_mpi_bgv_mnist_64.cpp;mpi_mnist_64.cpp;mpi_bgv_mnist_64.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
#   "helib;mpi;mkl")
# target_include_directories(test_mpi_bgv_mnist_64 PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(test_mpi_bgv_mnist_64 PUBLIC cxx_range_for)
# target_link_libraries(test_mpi_bgv_mnist_64 ${TORCH_LIBRARIES})

# MPI_BGV_64
create_target(
  test_mpi_bgv_mnist_64_test
  "test_mpi_bgv_mnist_64_test.cpp;mpi_mnist_64_test.cpp;mpi_bgv_mnist_64_test.cpp;emnist.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
  "helib;mpi;mkl")
target_include_directories(test_mpi_bgv_mnist_64_test PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_features(test_mpi_bgv_mnist_64_test PUBLIC cxx_range_for)
target_link_libraries(test_mpi_bgv_mnist_64_test ${TORCH_LIBRARIES})

# MPI_BGV_base
create_target(
  test_mpi_bgv_mnist_base
  "test_mpi_bgv_mnist_base.cpp;mpi_mnist_base.cpp;mpi_bgv_mnist_base.cpp;emnist.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
  "helib;mpi;mkl")
target_include_directories(test_mpi_bgv_mnist_base PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_features(test_mpi_bgv_mnist_base PUBLIC cxx_range_for)
target_link_libraries(test_mpi_bgv_mnist_base ${TORCH_LIBRARIES})

# MPI_BGV_base2
create_target(
  test_mpi_bgv_mnist_base2
  "test_mpi_bgv_mnist_base2.cpp;mpi_mnist_base2.cpp;mpi_bgv_mnist_base2.cpp;emnist.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
  "helib;mpi;mkl")
target_include_directories(test_mpi_bgv_mnist_base2 PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_features(test_mpi_bgv_mnist_base2 PUBLIC cxx_range_for)
target_link_libraries(test_mpi_bgv_mnist_base2 ${TORCH_LIBRARIES})

# MPI_BGV_2048
# create_target(
#   test_mpi_bgv_mnist_2048
#   "test_mpi_bgv_mnist_2048.cpp;mpi_mnist_2048.cpp;mpi_bgv_mnist_2048.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
#   "helib;mpi;mkl")
# target_include_directories(test_mpi_bgv_mnist_2048 PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(test_mpi_bgv_mnist_2048 PUBLIC cxx_range_for)
# target_link_libraries(test_mpi_bgv_mnist_2048 ${TORCH_LIBRARIES})

# MPI_BGV
# create_target(
#   test_mpi_bgv_mnist
#   "test_mpi_bgv_mnist.cpp;mpi_mnist.cpp;mpi_bgv_mnist.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
#   "helib;mpi;mkl")
# target_include_directories(test_mpi_bgv_mnist PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(test_mpi_bgv_mnist PUBLIC cxx_range_for)
# target_link_libraries(test_mpi_bgv_mnist ${TORCH_LIBRARIES})

# BGV
create_target(
  bgv_mnist
  "bgv_mnist.cpp;mnist_helper.cpp;bgv_mnist_helper.cpp;emnist.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
  "helib")
target_include_directories(bgv_mnist PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_features(bgv_mnist PUBLIC cxx_range_for)
target_link_libraries(bgv_mnist ${TORCH_LIBRARIES})

# BGV
create_target(
  bgv_mnist_shared
  "bgv_mnist_shared.cpp;mnist_helper_shared.cpp;bgv_mnist_helper_shared.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
  "helib")
target_include_directories(bgv_mnist_shared PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_features(bgv_mnist_shared PUBLIC cxx_range_for)
target_link_libraries(bgv_mnist_shared ${TORCH_LIBRARIES})

# create_target(
#   ckks_mnist
#   "ckks_mnist.cpp;mnist_helper.cpp;ckks_mnist_helper.cpp;${GENERAL_SRCS};emnist.cpp;${LINEAR_SRCS};${POWEROF2_SRCS}"
#   "helib")
# target_include_directories(ckks_mnist PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(ckks_mnist PUBLIC cxx_range_for)
# target_link_libraries(ckks_mnist ${TORCH_LIBRARIES})

# create_target(train_cifar10 "cifar10_train.cpp;cifar10.cpp" "")
# target_include_directories(train_cifar10 PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(train_cifar10 PUBLIC cxx_range_for)
# target_link_libraries(train_cifar10 ${TORCH_LIBRARIES})

# create_target(test_mpi_bgv_Cifar10_base_ptxt "test_mpi_bgv_Cifar10_base_ptxt.cpp;cifar10.cpp" "")
# target_include_directories(test_mpi_bgv_Cifar10_base_ptxt PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(test_mpi_bgv_Cifar10_base_ptxt PUBLIC cxx_range_for)
# target_link_libraries(test_mpi_bgv_Cifar10_base_ptxt ${TORCH_LIBRARIES})

# create_target(bgv_cifar10 "bgv_cifar10.cpp;cifar10.cpp" "")
# target_include_directories(bgv_cifar10 PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(bgv_cifar10 PUBLIC cxx_range_for)
# target_link_libraries(bgv_cifar10 ${TORCH_LIBRARIES})

if(MSVC)
  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
  add_custom_command(
    TARGET mnist
    POST_BUILD
    COMMAND ${CMAKE_COMMAND} -E copy_if_different ${TORCH_DLLS}
            $<TARGET_FILE_DIR:mnist>)
endif(MSVC)
